#!/usr/bin/python
# -*- coding: gbk -*-
'''
Created on 2010-7-7

@author: srcheng
'''

import datetime
import threading
import sys
from django.utils.encoding import smart_str
# xmpp
import client, protocol, debug

Debug = debug
Debug.DEBUGGING_IS_ON = 1
Debug.Debug.colors['gtalkbot'] = debug.color_green
DBG_GTALKBOT = 'gtalkbot'

# for debug
MainLog2File = False
MainDebugFlags = ['gtalkbot']
ClientLog2File = True
ClientDebugFlags = ['socket', 'roster']
OnlineOnly = True
INTERNAL_TEST = False

BOTNAME = 'GtalkBot'
BOTVER = 'v0.1'
BOTSERVER = 'talk.google.com'
BOTPORT = 5222

ERROR = 'ERROR'

# google capabilities
GC_GOOGLE_NOSAVE =  'google:nosave'
GC_ARCHIVE =        'http://jabber.org/protocol/archive'
GC_ARCHIVE_SAVE =   'http://jabber.org/protocol/archive#save'

class GtalkBot(object):
    def __init__(self, user, pwd, bot):
        global MainDebugFlags
        
        self._entryTime = datetime.datetime.now().strftime('%Y-%m-%d_%H-%M')
        
        self._bot = (bot == 'true')
        self._user = user
        self._pwd = pwd
        self._jid = protocol.JID(self._user)
        log = None
        if MainLog2File:
            if not MainDebugFlags: 
                MainDebugFlags=['always']
            log = '%s_%s_%s'%(self._user, self._entryTime, 'gtk.log')
            
        self._debugger = Debug.Debug(MainDebugFlags, log, time_stamp=1)
        self._debugger.debug_flags.append(DBG_GTALKBOT)
        self._receivers = {}
        self._receiversLock = threading.Lock()
        self._noSave = None
        self._archiveSave = None
        
    def Go(self):
        conres = self._connect()
        if not conres:
            self._debug('Connect to gtalk server failed:%s'%(conres), ERROR)
        else:
            self._debug('Connected via %s'%(conres))
            
        authres = self._auth()
        if not authres:
            self._debug('Auth failed:%s'%(authres), ERROR)
        else:
            self._debug('Authed wia %s'%(authres))
        
        self._init()
        self._registerHandlers()
        
        self._debug('Running...','main')
        self._goOn() 
               
    def _connect(self):
        global ClientDebugFlags
        
        log = None
        if ClientLog2File:
            if ClientDebugFlags:
                log = '%s_%s_%s'%(self._user, self._entryTime, 'xmpp.log')
                
        self._client = client.Client(server = self._jid.getDomain(), debug = ClientDebugFlags, log=log)
        self._debug('Connecting...', 'main')
        return self._client.connect(server = (BOTSERVER, BOTPORT))

    def _auth(self):        
        self._debug('Authenticating...','main')
        return self._client.auth(user=self._jid.getNode(), password=self._pwd)

    def _init(self):
        self._client.sendInitPresence()
        self._discoverCaps()
        self._roster = self._client.getRoster()
        if not OnlineOnly:
            self._receiversLock.acquire()
            for contact in self._roster.getItems():
                if self._roster.getSubscription(contact) in ['both']:
                    self._receivers[contact] = {}
                    self._debug('+++ %s added'%(contact), '_init')
            self._receiversLock.release()
        
    def _registerHandlers(self):
        self._client.RegisterHandler('message', self._messageCB)
        self._client.RegisterHandler('presence', self._presenceCB)
        self._client.RegisterHandler('iq', self._iqCB)
    
    def _debug(self, txt, severity = 'info'):
        self._debugger.Show(DBG_GTALKBOT, txt, severity)
        
    def _stepOn(self):
        try:
            self._client.Process(1)
        except KeyboardInterrupt:
            return 0
        return 1
    
    def _goOn(self):
        while self._stepOn(): 
            pass
        
    def _isOnline(self, typ):
        return typ in ['available', 'chat', 'away', 'online', None]
    
    def _discoverCaps(self):
        req = protocol.Iq(typ = 'get', queryNS = protocol.NS_DISCO_INFO, to=self._jid.getDomain())
        self._client.send(req)
        
    def _messageCB(self, conn, msg):
        frm = msg.getFrom().getStripped()
        typ = msg.getType()
        bdy = msg.getBody()
        id = msg.getID()
        
        if typ in ['error']:
            self._debug('invalid message type %s: %s.'%(frm, typ), 'msgCB')
            return
        if not bdy:
            self._debug('%s'%(msg), 'msgCB')
            return
        if len(bdy) > 20000:
            self._debug('message oversize!', ERROR)
            self._client.send(protocol.Message(frm, 'msg len should be less than 20000.'))
            return
        if self._bot:
            self._sendToOthers(frm, bdy, id)
        else:
            self._debug('%s: %s'%(frm, bdy), 'msgCB')
        
    def _sendToOthers(self, frm, bdy, id):
        self._debug('~~ %s: %s'%(frm, bdy), 'msgCB')
        name = protocol.JID(frm).getNode()
        self._receiversLock.acquire()
        contacts = self._receivers
        self._receiversLock.release()
        
        ''' debug '''
        if INTERNAL_TEST:
            tmpContacts = {}
            for contact in contacts:
                if contact in ['shirui.cheng@gmail.com', 'sanguoshamud@gmail.com']:
                    tmpContacts[contact] = contacts[contact]
            contacts = tmpContacts
        
        for contact in contacts:
            if contact != frm:
                self._sendToOne(contact, '<%s>: %s'%(name, bdy))
                
    def _sendToOne(self, jid, txt):
        msg = protocol.Message(to = jid, body = txt, typ = 'chat')
        ''' hack message for multicast '''
        msg = smart_str(msg).replace('<body>', 'None<body>')#.encode('utf-8')
        self._client.send(msg)
        
    def _presenceCB(self, conn, prs):
        jid = prs.getFrom()
        frm = jid.getStripped()
        res = jid.getResource()
        typ = prs.getType()
        if frm == self._user:
            self._debug('Ignore presence from myself.', 'prsCB')
            return
        if typ == 'subscribe':
            self._debug('%s authorized.'%(frm), 'prsCB')
            self._roster.Authorize(frm)
        elif OnlineOnly:
            if self._roster.getSubscription(frm) in ['both']:
                if self._isOnline(typ):
                    ''' add to contact list '''
                    if frm not in self._receivers:
                        self._receiversLock.acquire()
                        self._receivers[frm] = {}
                        self._receiversLock.release()
                    if res not in self._receivers[frm]:
                        self._receiversLock.acquire()
                        self._receivers[frm][res] = {}
                        self._receiversLock.release()
                        self._debug('+++ %s/%s added.'%(frm, res), 'prsCB')
                        
                else:
                    if frm in self._receivers:
                        ''' remove from contact list '''
                        if res in self._receivers[frm]:
                            self._receiversLock.acquire()
                            del self._receivers[frm][res]
                            self._receiversLock.release()
                            self._debug('--- %s/%s deleted.'%(frm, res), 'prsCB')
                
    def _iqCB(self, conn, iq):
        ns = iq.getQueryNS()
        frm = iq.getFrom()
        typ = iq.getType()
        self._debug('iq: from(%s), type(%s), ns(%s)'%(frm, typ, ns), 'iqCB')
        if ns == protocol.NS_DISCO_INFO:
            self._handleDiscoInfo(conn, iq)
        else: # TODO
            pass
                
    def _handleDiscoInfo(self, conn, iq):
        for item in iq.getTag('query').getTags('feature'):
            if item.getAttr('var') == GC_GOOGLE_NOSAVE:
                self._noSave = True
            elif item.getAttr('var') == GC_ARCHIVE_SAVE:
                self._archiveSave = True
                
# main: test
if __name__ == "__main__":
    if len(sys.argv) == 4:
        user = sys.argv[1]
        pwd = sys.argv[2]
        bot = sys.argv[3].lower()
        print 'GtalkBot: %s(%s)'%(user, bot)
        myBot = GtalkBot(user = user, pwd = pwd, bot = bot)
        myBot.Go()